from utils.utils import *
from config.config import *
import joblib
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import os

import numpy as np
import pandas as pd
def feature_type(feat):
    between = feat.split("_", 1)[1].split("(", 1)[0]  # 取 "_" 和 "(" 之间的内容
    if between == "M":
        return "mean"
    elif between == "S":
        return "std"
    
    return None
import re


def compute_signature_from_roi(roi_df, coef_dict):

    N = roi_df.shape[0]
    roi_contrib = pd.Series(0.0, index=roi_df.index)
    for feat, beta in coef_dict.items():
        ftype = feature_type(feat)
        raw_feat = re.search(r"\((.*?)\)", feat).group(1)

        if raw_feat not in roi_df.columns:
            return 

        # 区分来源（Region / Collagen / Nuclei）
        if feat.startswith("Region"):
            weights = None
        else:  # Collagen / Nuclei
            weights = roi_df["Saliency.SaliencyScore"].values
            weights = np.clip(weights, 1e-8, None)  # 避免0

        x = roi_df[raw_feat].replace([np.inf, -np.inf], np.nan).fillna(0).values
        N = len(x)
        if ftype == "mean":
            if weights is None:  # Region: 均值
                contribs = (beta / N) * x
            else:  # 加权均值
                contribs = beta * (weights / weights.sum()) * x

            roi_contrib += contribs

        elif ftype == "std":

            if weights is None:  # Region: std
                mu = np.nanmean(x)
                std = np.nanstd(x)
                if std > 0:
                    contribs = beta * ((x - mu) ** 2) / (N * std)
                    roi_contrib += contribs
            else:  # 加权 std
                mu = np.average(x, weights=weights)
                var = np.average((x - mu) ** 2, weights=weights)
                std = np.sqrt(var)
                if std > 0:
                    contribs = beta * (weights * (x - mu) ** 2) / (weights.sum() * std)
                    roi_contrib += contribs

    result = roi_contrib.to_frame()
    result.columns = [f"{s.replace('Roi','Nuclei')}_score"]
    return result



s_list = ['Collagen','Roi']
save_path = opj(base_path,'Mutils','output','HiPS','BCNB','cTMEfeats')

for s in s_list:
    columns = pd.read_csv(opj(lasso_feature_path,str(split_random_state_list[0]),str(selectK_random_state_list[0]),f'{s}.csv')).drop(columns=exclude_columns).columns.tolist()
    columns_processed = pd.Series(columns).str.extract(r'\((.*?)\)')[0].tolist()
    filename = 'perSlideCollagenFeatures'
    if s == 'Roi':
        filename = 'perSlideROISummaries'
    roi_save_path = opj(save_path,f"perROI{s.replace('Roi','Nuclei')}Distribution")
    md(roi_save_path)
    temp_path = opj(base_path,'Mutils','output','HiPS','BCNB','cTMEfeats',filename)
    data_list = [p.replace('.csv','') for p in ol(temp_path)]
    data_list.sort(key=int)
    for c in data_list:
        roi_level_data = pd.read_csv(opj(temp_path,f'{c}.csv'),index_col=0)
        model = joblib.load(opj(signature_model_weight_path,s.replace('Roi','Nuclei'),'model.pkl')).get_model()
        coef =  dict(zip(columns, model.coef_[0]))
        result = compute_signature_from_roi(roi_level_data,coef)
        if result is not None:
            result.to_csv(opj(roi_save_path,f'{c}.csv'))